{-# LANGUAGE TypeSynonymInstances       #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE EmptyDataDecls             #-}
{-# LANGUAGE TypeFamilies               #-}
{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  Orca.Language.Prim
-- Copyright   :  (c) Stephen Tetley 2013
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Typed layer to build code.
--
--------------------------------------------------------------------------------

module Orca.Language.Prim
  (

    PExpr(..)       
  , IString
  , IBool
  , IInt
  , IFloat
  , TableIx
  , tableIx

  , IRate(..)
  , KRate(..)
  , ARate(..)

  , Sig    
  , ISig
  , KSig
  , ASig

  , InitValue
  , rateId
  , IKA_Signal
  , IK_Signal
  , KA_Signal

  , unaryOp
  , binOp
  , relOp

  , GenSig
  , GenISig
  , GenKSig
  , GenASig
  , GenSig2
  , GenSig3
  , GenSig4


  , Instr
  , InstCode
  , InstDef

  , globalAssignTable
  , globalAssign        -- only exported for defintion of zakinit like funs
  , globalInvoke        -- -ditto-
  , instr

  , definstr
  , defvar
  , instbody

  , isignal
  , ksignal
  , asignal

  , isignal2
  , ksignal2
  , asignal2

  , liftS

  , signalI
  , signalK
  , signalA

  , (&=)
  , (=%=)

  , opcode0
  , opcode1
  , opcode2
  , opcode3
  , opcode4

  , function
  , pfield

  , SMax
  , (.+.)  
  , (.-.)
  , (.*.)
  , (./.)


  , CastA
  , castA
  , CastK
  , castK

  , (*^)
  , (^*)
  , (/^)
  , (^/)
  , (+^)
  , (^+)
  , (-^)
  , (^-)

  , cond
  , (.==.)
  , (./=.)
  , (.>.)
  , (.<.) 
  , (.>=.)
  , (.<=.) 
    
  ) where


import Orca.Prim.InstMon
import Orca.Prim.OrchMon
import Orca.Prim.Syntax


import Control.Applicative
import Data.String


-- | Phantom type wrapping Expressions.
--
newtype PExpr a = PExpr { getPExpr :: Expr } 



type IString        = PExpr String
type IBool          = PExpr Bool
type IInt           = PExpr Int
type IFloat         = PExpr Float


-- | Separate type for table indices (with are otherwise Int) 
-- disallowing Int operations e.g. Eq, Ord, Num.
--
-- This is analogue to disallowing pointer arithmetic.
-- 
data TableIxTag

type TableIx        = PExpr TableIxTag

tableIx :: Int -> TableIx
tableIx = PExpr . Value . Int


data IRate          = IRate deriving (Show)
data KRate          = KRate deriving (Show)
data ARate          = ARate deriving (Show)


-- (Export Sig rather than PExpr)
type Sig a          = PExpr a


type ISig           = Sig IRate
type KSig           = Sig KRate
type ASig           = Sig ARate



class InitValue ix where
   rateId :: ix -> RateId

instance InitValue ARate       where rateId _ = A_RATE
instance InitValue KRate       where rateId _ = K_RATE
instance InitValue IRate       where rateId _ = I_RATE


class InitValue rate => IKA_Signal rate 

instance IKA_Signal ARate
instance IKA_Signal KRate
instance IKA_Signal IRate

class IK_Signal a

instance IK_Signal KRate
instance IK_Signal IRate

class KA_Signal a

instance KA_Signal KRate
instance KA_Signal ARate



instance IsString IString where
  fromString = PExpr . Value . String


instance Num IInt where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs _     = error $ "abs - no implementation for VInt"
  signum _  = error $ "signum - no implementation for VInt"
  fromInteger = PExpr . Value . Int . fromInteger


instance Num IFloat where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs _     = error $ "abs - no implementation for VFloat"
  signum _  = error $ "signum - no implementation for VFloat"
  fromInteger = PExpr . Value . Float . fromInteger

instance Num ISig where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs _     = error $ "abs - no implementation for ISig"
  signum _  = error $ "signum - no implementation for ISig"
  fromInteger = PExpr . Value . Float . fromInteger


instance Num KSig where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs _     = error $ "abs - no implementation for KSig"
  signum _  = error $ "signum - no implementation for KSig"
  fromInteger = PExpr . Value . Float . fromInteger


instance Num ASig where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs _     = error $ "abs - no implementation for ASig"
  signum _  = error $ "signum - no implementation for ASig"
  fromInteger = PExpr . Value . Float . fromInteger


instance Fractional IFloat where
  (/)       = binOp DIV
  recip _   = error $ "recip - no implementation for VFloat"
  fromRational = PExpr . Value . Float . fromRational


instance Fractional ISig where
  (/)       = binOp DIV
  recip _   = error $ "recip - no implementation for ISig"
  fromRational = PExpr . Value . Float . fromRational


instance Fractional KSig where
  (/)       = binOp DIV
  recip _   = error $ "recip - no implementation for KSig"
  fromRational = PExpr . Value . Float . fromRational

instance Fractional ASig where
  (/)       = binOp DIV
  recip _   = error $ "recip - no implementation for ASig"
  fromRational = PExpr . Value . Float . fromRational


-- Can define NUM instances for empty types as PExpr is a 
-- Phantom.

unaryOp :: UnaryOp -> PExpr a -> PExpr a
unaryOp op a = PExpr $ 
    UnaryE op (getPExpr a)




binOp :: BinOp -> PExpr a -> PExpr a -> PExpr a
binOp op a b = PExpr $ 
    BinaryE op (getPExpr a) (getPExpr b)


relOp :: RelOp -> PExpr a -> PExpr a -> PExpr Bool
relOp op a b = PExpr $ 
    RelE op (getPExpr a) (getPExpr b)




-- | The return type of a opcode is not actually a signal, it
-- is /code that generates a signal/.
--
-- This allows opcalls to be used both to generate a signal 
-- (creating a fresh var) and to reassign a signal.
--
newtype GenSig rate = GenSig { getGenSig :: Expr2 }

type GenISig = GenSig IRate
type GenKSig = GenSig KRate
type GenASig = GenSig ARate


type GenSig2 a = GenSig (a,a)
type GenSig3 a = GenSig (a,a,a)
type GenSig4 a = GenSig (a,a,a,a)



--------------------------------------------------------------------------------



newtype Instr = Instr { getInstr :: Int -> Either ErrMsg Instrument }


newtype InstCode a = InstCode { getInstCode :: Code a }

type InstDef = InstCode ()



instance Functor InstCode where 
  fmap f ma = InstCode $ getInstCode ma >>= \a -> return (f a)

instance Applicative InstCode where 
  pure a    = InstCode $ return a
  mf <*> ma = InstCode $ getInstCode mf <*> getInstCode ma

instance Monad InstCode where
  return   = pure
  ma >>= k = InstCode $ getInstCode ma >>= \a -> getInstCode (k a)
                       



globalAssignTable :: String -> RateId -> PrimId -> [Expr] -> Orch TableIx
globalAssignTable ss ri pid exprs = 
    tellGlobalStmt stmt >> return (PExpr $ Value $ Var vid)
  where
    vid   = Global ri ss
    stmt  = GlobalDecl vid pid exprs

globalAssign :: String -> RateId -> PrimId -> [Expr] -> Orch (Sig a)
globalAssign ss ri pid exprs = 
    tellGlobalStmt stmt >> return (PExpr $ Value $ Var vid)
  where
    vid   = Global ri ss
    stmt  = GlobalDecl vid pid exprs

globalInvoke :: PrimId -> [Expr] -> Orch ()
globalInvoke pid exprs = tellGlobalStmt stmt >> return ()
  where
    stmt = GlobalOpCall pid exprs


instr :: Int -> Instr -> Orch ()
instr i def = either throwErr tellInstrument $ getInstr def i



defvar :: String -> PExpr a -> InstCode (PExpr a)
defvar ss val = InstCode $ do
    tellPureInit $ LocalDecl vid (getPExpr val)
    return $ PExpr $ Value $ Var vid
  where
    vid = DeclVar I_RATE ss 
    
       



definstr :: InstDef -> Instr
definstr ma = Instr $ \i ->  
    sk i <$> execCode (getInstCode ma)
  where
    sk i (decls, stmts) = Instrument { instr_num   = i
                                     , instr_decls = decls                                    
                                     , instr_body  = stmts }


instbody :: Code a -> InstCode ()
instbody ma = InstCode $ ma >> return ()

    



isignal :: GenSig IRate -> Code (Sig IRate)
isignal gexpr = do
    vid     <- newVarId I_RATE
    tellStmt $ SigDef1 vid (getGenSig gexpr)
    return $ PExpr $ Value $ Var vid

ksignal :: GenSig KRate -> Code (Sig KRate)
ksignal gexpr = do
    vid     <- newVarId K_RATE
    tellStmt $ SigDef1 vid (getGenSig gexpr)
    return $ PExpr $ Value $ Var vid

asignal :: GenSig ARate -> Code (Sig ARate)
asignal gexpr = do
    vid     <- newVarId A_RATE
    tellStmt $ SigDef1 vid (getGenSig gexpr)
    return $ PExpr $ Value $ Var vid

isignal2 :: GenSig2 IRate -> Code (Sig IRate, Sig IRate)
isignal2 gexpr = do
    v1    <- newVarId I_RATE
    v2    <- newVarId I_RATE
    tellStmt $ SigDef2 v1 v2 (getGenSig gexpr)
    return $ (PExpr $ Value $ Var v1, PExpr $ Value $ Var v2)

ksignal2 :: GenSig2 KRate -> Code (Sig KRate, Sig KRate)
ksignal2 gexpr = do
    v1    <- newVarId K_RATE
    v2    <- newVarId K_RATE
    tellStmt $ SigDef2 v1 v2 (getGenSig gexpr)
    return $ (PExpr $ Value $ Var v1, PExpr $ Value $ Var v2)

asignal2 :: GenSig2 ARate -> Code (Sig ARate, Sig ARate)
asignal2 gexpr = do
    v1    <- newVarId A_RATE
    v2    <- newVarId A_RATE
    tellStmt $ SigDef2 v1 v2 (getGenSig gexpr)
    return $ (PExpr $ Value $ Var v1, PExpr $ Value $ Var v2)


liftS :: Sig rate -> GenSig rate
liftS (PExpr a) = GenSig $ Pure a


signalI :: IFloat -> Sig IRate
signalI = PExpr . getPExpr

signalK :: IFloat -> Sig KRate
signalK = castK . signalI

signalA :: IFloat -> Sig ARate
signalA = castA . signalI


infixr 1 &=

-- | Assignment
--
(&=) :: Sig rate -> GenSig rate -> Code ()
(&=) (PExpr val) ma = case val of
    Value (Var vid) -> let code = getGenSig ma 
                       in tellStmt $ Update vid code
    _          -> throwErr "Assignment - lhs not a signal (lvar)."


infixr 1 =%=

-- | Assignment
--
(=%=) :: Sig rate -> Sig rate -> Code ()
(=%=) (PExpr val) expr = case val of
    Value (Var vid) -> let code = Pure $ getPExpr expr 
                       in tellStmt $ Update vid code
    _          -> throwErr "Assignment - lhs not a signal (lvar)."




-- Opcodes return code, rather than /tell/ code. 
-- The exception is opcode0 e.g. out, outs
--
opcode0 :: String -> [Expr] -> Code ()
opcode0 ss vs = tellStmt $ OpCallS (primId ss) vs


opcode1 :: String -> [Expr] -> GenSig rate
opcode1 ss vs = GenSig $ OpCallE (primId ss) vs

opcode2 :: String -> [Expr] -> GenSig2 rate
opcode2 ss vs = GenSig $ OpCallE (primId ss) vs

opcode3 :: String -> [Expr] -> GenSig3 rate
opcode3 ss vs = GenSig $ OpCallE (primId ss) vs

opcode4 :: String -> [Expr] -> GenSig4 rate
opcode4 ss vs = GenSig $ OpCallE (primId ss) vs



function :: String -> [Expr] -> PExpr ans
function name vs = PExpr $ FunCall (primId name) vs

pfield :: Int -> PExpr a
pfield = PExpr . Value . Pfield

type family SMax a b :: *

type instance SMax ARate ARate = ARate
type instance SMax ARate KRate = ARate
type instance SMax KRate ARate = ARate
type instance SMax ARate IRate = ARate
type instance SMax IRate ARate = ARate

type instance SMax KRate KRate = KRate
type instance SMax KRate IRate = KRate
type instance SMax IRate KRate = KRate

type instance SMax IRate IRate = IRate


infixl 6 .+.

(.+.) :: c ~ SMax a b => Sig a -> Sig b -> Sig c
(.+.) s1 s2 = PExpr $ BinaryE ADD (getPExpr s1) (getPExpr s2)


infixl 6 .-.

(.-.) :: c ~ SMax a b => Sig a -> Sig b -> Sig c
(.-.) s1 s2 = PExpr $ BinaryE SUB (getPExpr s1) (getPExpr s2)

infixl 7 .*.

(.*.) :: c ~ SMax a b => Sig a -> Sig b -> Sig c
(.*.) s1 s2 = PExpr $ BinaryE MUL (getPExpr s1) (getPExpr s2)

infixl 7 ./.

(./.) :: c ~ SMax a b => Sig a -> Sig b -> Sig c
(./.) s1 s2 = PExpr $ BinaryE DIV (getPExpr s1) (getPExpr s2)

class CastK a where castK :: Sig a -> Sig KRate

instance CastK IRate where
  castK = PExpr . getPExpr

instance CastK KRate where
  castK = id


class CastA a where castA :: Sig a -> Sig ARate

instance CastA IRate where
  castA = PExpr . getPExpr

instance CastA KRate where
  castA = PExpr . getPExpr

instance CastA ARate where
  castA = id



infixr 7 *^
infixl 7 ^*

-- | Multiply a signal by an I-rate constant.
--
(*^) :: (IKA_Signal a, Num (Sig a)) => IFloat -> Sig a -> Sig a
(*^) d s = (cast d :: Sig a) * s
  where 
    cast = PExpr . getPExpr

(^*) :: (IKA_Signal a, Num (Sig a)) => Sig a -> IFloat -> Sig a
(^*) s d = s * (cast d :: Sig a) 
  where 
    cast = PExpr . getPExpr



infixr 7 /^
infixl 7 ^/
-- | Divide a signal by an I-rate constant.
--
(/^) :: (IKA_Signal a, Fractional (Sig a)) => IFloat -> Sig a -> Sig a
(/^) d s = (cast d :: Sig a) / s
  where 
    cast = PExpr . getPExpr

(^/) :: (IKA_Signal a, Fractional (Sig a)) => Sig a -> IFloat -> Sig a
(^/) s d = s / (cast d :: Sig a)
  where
    cast = PExpr . getPExpr


infixr 6 +^
infixl 6 ^+

-- | Addition of a signal and an I-rate constant.
--
(+^) :: (IKA_Signal a, Num (Sig a)) => IFloat -> Sig a -> Sig a
(+^) d s = (cast d :: Sig a) + s
  where 
    cast = PExpr . getPExpr

(^+) :: (IKA_Signal a, Num (Sig a)) => Sig a -> IFloat -> Sig a
(^+) s d = s + (cast d :: Sig a) 
  where
    cast = PExpr . getPExpr


infixr 6 -^
infixl 6 ^-

-- | Subtraction from a signal by an I-rate constant.
--
(-^) :: (IKA_Signal a, Num (Sig a)) => IFloat -> Sig a -> Sig a
(-^) d s = (cast d :: Sig a) - s
  where 
    cast = PExpr . getPExpr



(^-) :: (IKA_Signal a, Num (Sig a)) => Sig a -> IFloat -> Sig a
(^-) s d = s - (cast d :: Sig a)
  where 
    cast = PExpr . getPExpr


cond :: IBool -> PExpr a -> PExpr a -> PExpr a
cond test a b = PExpr $ CondE (getPExpr test) (getPExpr a) (getPExpr b)


infix 4 .==.

(.==.) :: Eq a => PExpr a -> PExpr a -> IBool
(.==.) = relOp EQU


infix 4 ./=. 

(./=.) :: Eq a => PExpr a -> PExpr a -> IBool
(./=.) = relOp NEQU


infix 4 .>. 

(.>.) :: Eq a => PExpr a -> PExpr a -> IBool
(.>.) = relOp GREATTHAN

infix 4 .<. 

(.<.) :: Eq a => PExpr a -> PExpr a -> IBool
(.<.) = relOp LESSTHAN

infix 4 .>=. 

(.>=.) :: Eq a => PExpr a -> PExpr a -> IBool
(.>=.) = relOp GTEQU

infix 4 .<=. 

(.<=.) :: Eq a => PExpr a -> PExpr a -> IBool
(.<=.) = relOp LTEQU